/*
 * Copyright 2024-2024 the original author or authors.
 */

package io.modelcontextprotocol.server;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;

import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.util.Assert;
import io.modelcontextprotocol.util.Utils;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

/**
 * MCP server features specification that a particular server can choose to support.
 *
 * @author Dariusz Jędrzejczyk
 * @author Jihoon Kim
 */
public class McpServerFeatures {

	/**
	 * Asynchronous server features specification.
	 */
	public static class Async {
		private final McpSchema.Implementation serverInfo;
		private final McpSchema.ServerCapabilities serverCapabilities;
		private final List<AsyncToolSpecification> tools;
		private final Map<String, AsyncResourceSpecification> resources;
		private final List<McpSchema.ResourceTemplate> resourceTemplates;
		private final Map<String, AsyncPromptSpecification> prompts;
		private final Map<McpSchema.CompleteReference, AsyncCompletionSpecification> completions;
		private final List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootsChangeConsumers;
		private final String instructions;

		/**
		 * Create an instance and validate the arguments.
		 * @param serverInfo The server implementation details
		 * @param serverCapabilities The server capabilities
		 * @param tools The list of tool specifications
		 * @param resources The map of resource specifications
		 * @param resourceTemplates The list of resource templates
		 * @param prompts The map of prompt specifications
		 * @param rootsChangeConsumers The list of consumers that will be notified when
		 * the roots list changes
		 * @param instructions The server instructions text
		 */
		public Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities,
					 List<AsyncToolSpecification> tools, Map<String, AsyncResourceSpecification> resources,
					 List<McpSchema.ResourceTemplate> resourceTemplates,
					 Map<String, AsyncPromptSpecification> prompts,
					 Map<McpSchema.CompleteReference, AsyncCompletionSpecification> completions,
					 List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootsChangeConsumers,
					 String instructions) {

			Assert.notNull(serverInfo, "Server info must not be null");

			this.serverInfo = serverInfo;
			this.serverCapabilities = (serverCapabilities != null) ? serverCapabilities
					: new McpSchema.ServerCapabilities(null, // completions
					null, // experimental
					new McpSchema.ServerCapabilities.LoggingCapabilities(), // Enable
					// logging
					// by
					// default
					!Utils.isEmpty(prompts) ? new McpSchema.ServerCapabilities.PromptCapabilities(false) : null,
					!Utils.isEmpty(resources)
							? new McpSchema.ServerCapabilities.ResourceCapabilities(false, false) : null,
					!Utils.isEmpty(tools) ? new McpSchema.ServerCapabilities.ToolCapabilities(false) : null);

			this.tools = (tools != null) ? tools : new ArrayList<>();
			this.resources = (resources != null) ? resources : new HashMap<>();
			this.resourceTemplates = (resourceTemplates != null) ? resourceTemplates : new ArrayList<>();
			this.prompts = (prompts != null) ? prompts : new HashMap<>();
			this.completions = (completions != null) ? completions : new HashMap<>();
			this.rootsChangeConsumers = (rootsChangeConsumers != null) ? rootsChangeConsumers : new ArrayList<>();
			this.instructions = instructions;
		}

		public McpSchema.Implementation serverInfo() {
			return serverInfo;
		}

		public McpSchema.ServerCapabilities serverCapabilities() {
			return serverCapabilities;
		}

		public List<AsyncToolSpecification> tools() {
			return tools;
		}

		public Map<String, AsyncResourceSpecification> resources() {
			return resources;
		}

		public List<McpSchema.ResourceTemplate> resourceTemplates() {
			return resourceTemplates;
		}

		public Map<String, AsyncPromptSpecification> prompts() {
			return prompts;
		}

		public Map<McpSchema.CompleteReference, AsyncCompletionSpecification> completions() {
			return completions;
		}

		public List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootsChangeConsumers() {
			return rootsChangeConsumers;
		}

		public String instructions() {
			return instructions;
		}

		/**
		 * Convert a synchronous specification into an asynchronous one and provide
		 * blocking code offloading to prevent accidental blocking of the non-blocking
		 * transport.
		 * @param syncSpec a potentially blocking, synchronous specification.
		 * @return a specification which is protected from blocking calls specified by the
		 * user.
		 */
		public static Async fromSync(Sync syncSpec) {
			List<AsyncToolSpecification> tools = new ArrayList<>();
			for (SyncToolSpecification tool : syncSpec.getTools()) {
				tools.add(AsyncToolSpecification.fromSync(tool));
			}

			Map<String, AsyncResourceSpecification> resources = new HashMap<>();
			for (Map.Entry<String, SyncResourceSpecification> entry : syncSpec.getResources().entrySet()) {
				resources.put(entry.getKey(), AsyncResourceSpecification.fromSync(entry.getValue()));
			}

			Map<String, AsyncPromptSpecification> prompts = new HashMap<>();
			for (Map.Entry<String, SyncPromptSpecification> entry : syncSpec.getPrompts().entrySet()) {
				prompts.put(entry.getKey(), AsyncPromptSpecification.fromSync(entry.getValue()));
			}

			Map<McpSchema.CompleteReference, AsyncCompletionSpecification> completions = new HashMap<>();
			for (Map.Entry<McpSchema.CompleteReference, SyncCompletionSpecification> entry : syncSpec.getCompletions().entrySet()) {
				completions.put(entry.getKey(), AsyncCompletionSpecification.fromSync(entry.getValue()));
			}

			List<BiFunction<McpAsyncServerExchange, List<McpSchema.Root>, Mono<Void>>> rootChangeConsumers = new ArrayList<>();
			for (BiConsumer<McpSyncServerExchange, List<McpSchema.Root>> rootChangeConsumer : syncSpec.getRootsChangeConsumers()) {
				rootChangeConsumers.add((exchange, list) -> Mono
						.<Void>fromRunnable(() -> rootChangeConsumer.accept(new McpSyncServerExchange(exchange), list))
						.subscribeOn(Schedulers.boundedElastic()));
			}

			return new Async(syncSpec.getServerInfo(), syncSpec.getServerCapabilities(), tools, resources,
					syncSpec.getResourceTemplates(), prompts, completions, rootChangeConsumers, syncSpec.getInstructions());
		}
	}

	/**
	 * Synchronous server features specification.
	 */
	public static class Sync {
		private final McpSchema.Implementation serverInfo;
		private final McpSchema.ServerCapabilities serverCapabilities;
		private final List<SyncToolSpecification> tools;
		private final Map<String, SyncResourceSpecification> resources;
		private final List<McpSchema.ResourceTemplate> resourceTemplates;
		private final Map<String, SyncPromptSpecification> prompts;
		private final Map<McpSchema.CompleteReference, SyncCompletionSpecification> completions;
		private final List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> rootsChangeConsumers;
		private final String instructions;

		/**
		 * Create an instance and validate the arguments.
		 * @param serverInfo The server implementation details
		 * @param serverCapabilities The server capabilities
		 * @param tools The list of tool specifications
		 * @param resources The map of resource specifications
		 * @param resourceTemplates The list of resource templates
		 * @param prompts The map of prompt specifications
		 * @param rootsChangeConsumers The list of consumers that will be notified when
		 * the roots list changes
		 * @param instructions The server instructions text
		 */
		public Sync(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities serverCapabilities,
					List<SyncToolSpecification> tools, Map<String, SyncResourceSpecification> resources,
					List<McpSchema.ResourceTemplate> resourceTemplates,
					Map<String, SyncPromptSpecification> prompts,
					Map<McpSchema.CompleteReference, SyncCompletionSpecification> completions,
					List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> rootsChangeConsumers,
					String instructions) {

			Assert.notNull(serverInfo, "Server info must not be null");

			this.serverInfo = serverInfo;
			this.serverCapabilities = (serverCapabilities != null) ? serverCapabilities
					: new McpSchema.ServerCapabilities(null, // completions
					null, // experimental
					new McpSchema.ServerCapabilities.LoggingCapabilities(), // Enable
					// logging
					// by
					// default
					!Utils.isEmpty(prompts) ? new McpSchema.ServerCapabilities.PromptCapabilities(false) : null,
					!Utils.isEmpty(resources)
							? new McpSchema.ServerCapabilities.ResourceCapabilities(false, false) : null,
					!Utils.isEmpty(tools) ? new McpSchema.ServerCapabilities.ToolCapabilities(false) : null);

			this.tools = (tools != null) ? tools : new ArrayList<>();
			this.resources = (resources != null) ? resources : new HashMap<>();
			this.resourceTemplates = (resourceTemplates != null) ? resourceTemplates : new ArrayList<>();
			this.prompts = (prompts != null) ? prompts : new HashMap<>();
			this.completions = (completions != null) ? completions : new HashMap<>();
			this.rootsChangeConsumers = (rootsChangeConsumers != null) ? rootsChangeConsumers : new ArrayList<>();
			this.instructions = instructions;
		}

		public McpSchema.Implementation getServerInfo() {
			return serverInfo;
		}

		public McpSchema.ServerCapabilities getServerCapabilities() {
			return serverCapabilities;
		}

		public List<SyncToolSpecification> getTools() {
			return tools;
		}

		public Map<String, SyncResourceSpecification> getResources() {
			return resources;
		}

		public List<McpSchema.ResourceTemplate> getResourceTemplates() {
			return resourceTemplates;
		}

		public Map<String, SyncPromptSpecification> getPrompts() {
			return prompts;
		}

		public Map<McpSchema.CompleteReference, SyncCompletionSpecification> getCompletions() {
			return completions;
		}

		public List<BiConsumer<McpSyncServerExchange, List<McpSchema.Root>>> getRootsChangeConsumers() {
			return rootsChangeConsumers;
		}

		public String getInstructions() {
			return instructions;
		}
	}

	/**
	 * Specification of a tool with its asynchronous handler function.
	 */
	public static class AsyncToolSpecification {
		private final McpSchema.Tool tool;
		private final BiFunction<McpAsyncServerExchange, Map<String, Object>, Mono<McpSchema.CallToolResult>> call;

		public AsyncToolSpecification(McpSchema.Tool tool,
									  BiFunction<McpAsyncServerExchange, Map<String, Object>, Mono<McpSchema.CallToolResult>> call) {
			this.tool = tool;
			this.call = call;
		}

		public McpSchema.Tool tool() {
			return tool;
		}

		public BiFunction<McpAsyncServerExchange, Map<String, Object>, Mono<McpSchema.CallToolResult>> call() {
			return call;
		}

		public static AsyncToolSpecification fromSync(SyncToolSpecification tool) {
			// FIXME: This is temporary, proper validation should be implemented
			if (tool == null) {
				return null;
			}
			return new AsyncToolSpecification(tool.tool(),
					(exchange, map) -> Mono
							.fromCallable(() -> tool.call().apply(new McpSyncServerExchange(exchange), map))
							.subscribeOn(Schedulers.boundedElastic()));
		}
	}

	/**
	 * Specification of a resource with its asynchronous handler function.
	 */
	public static class AsyncResourceSpecification {
		private final McpSchema.Resource resource;
		private final BiFunction<McpAsyncServerExchange, McpSchema.ReadResourceRequest, Mono<McpSchema.ReadResourceResult>> readHandler;

		public AsyncResourceSpecification(McpSchema.Resource resource,
										  BiFunction<McpAsyncServerExchange, McpSchema.ReadResourceRequest, Mono<McpSchema.ReadResourceResult>> readHandler) {
			this.resource = resource;
			this.readHandler = readHandler;
		}

		public McpSchema.Resource resource() {
			return resource;
		}

		public BiFunction<McpAsyncServerExchange, McpSchema.ReadResourceRequest, Mono<McpSchema.ReadResourceResult>> readHandler() {
			return readHandler;
		}

		public static AsyncResourceSpecification fromSync(SyncResourceSpecification resource) {
			// FIXME: This is temporary, proper validation should be implemented
			if (resource == null) {
				return null;
			}
			return new AsyncResourceSpecification(resource.resource(),
					(exchange, req) -> Mono
							.fromCallable(() -> resource.readHandler().apply(new McpSyncServerExchange(exchange), req))
							.subscribeOn(Schedulers.boundedElastic()));
		}
	}

	/**
	 * Specification of a prompt template with its asynchronous handler function.
	 */
	public static class AsyncPromptSpecification {
		private final McpSchema.Prompt prompt;
		private final BiFunction<McpAsyncServerExchange, McpSchema.GetPromptRequest, Mono<McpSchema.GetPromptResult>> promptHandler;

		public AsyncPromptSpecification(McpSchema.Prompt prompt,
										BiFunction<McpAsyncServerExchange, McpSchema.GetPromptRequest, Mono<McpSchema.GetPromptResult>> promptHandler) {
			this.prompt = prompt;
			this.promptHandler = promptHandler;
		}

		public McpSchema.Prompt prompt() {
			return prompt;
		}

		public BiFunction<McpAsyncServerExchange, McpSchema.GetPromptRequest, Mono<McpSchema.GetPromptResult>> promptHandler() {
			return promptHandler;
		}

		public static AsyncPromptSpecification fromSync(SyncPromptSpecification prompt) {
			// FIXME: This is temporary, proper validation should be implemented
			if (prompt == null) {
				return null;
			}
			return new AsyncPromptSpecification(prompt.prompt(),
					(exchange, req) -> Mono
							.fromCallable(() -> prompt.promptHandler().apply(new McpSyncServerExchange(exchange), req))
							.subscribeOn(Schedulers.boundedElastic()));
		}
	}

	/**
	 * Specification of a completion handler function with asynchronous execution support.
	 */
	public static class AsyncCompletionSpecification {
		private final McpSchema.CompleteReference referenceKey;
		private final BiFunction<McpAsyncServerExchange, McpSchema.CompleteRequest, Mono<McpSchema.CompleteResult>> completionHandler;

		public AsyncCompletionSpecification(McpSchema.CompleteReference referenceKey,
											BiFunction<McpAsyncServerExchange, McpSchema.CompleteRequest, Mono<McpSchema.CompleteResult>> completionHandler) {
			this.referenceKey = referenceKey;
			this.completionHandler = completionHandler;
		}

		public McpSchema.CompleteReference referenceKey() {
			return referenceKey;
		}

		public BiFunction<McpAsyncServerExchange, McpSchema.CompleteRequest, Mono<McpSchema.CompleteResult>> completionHandler() {
			return completionHandler;
		}

		public static AsyncCompletionSpecification fromSync(SyncCompletionSpecification completion) {
			if (completion == null) {
				return null;
			}
			return new AsyncCompletionSpecification(completion.referenceKey(),
					(exchange, request) -> Mono.fromCallable(
									() -> completion.completionHandler().apply(new McpSyncServerExchange(exchange), request))
							.subscribeOn(Schedulers.boundedElastic()));
		}
	}

	/**
	 * Specification of a tool with its synchronous handler function.
	 */
	public static class SyncToolSpecification {
		private final McpSchema.Tool tool;
		private final BiFunction<McpSyncServerExchange, Map<String, Object>, McpSchema.CallToolResult> call;

		public SyncToolSpecification(McpSchema.Tool tool,
									 BiFunction<McpSyncServerExchange, Map<String, Object>, McpSchema.CallToolResult> call) {
			this.tool = tool;
			this.call = call;
		}

		public McpSchema.Tool tool() {
			return tool;
		}

		public BiFunction<McpSyncServerExchange, Map<String, Object>, McpSchema.CallToolResult> call() {
			return call;
		}
	}

	/**
	 * Specification of a resource with its synchronous handler function.
	 */
	public static class SyncResourceSpecification {
		private final McpSchema.Resource resource;
		private final BiFunction<McpSyncServerExchange, McpSchema.ReadResourceRequest, McpSchema.ReadResourceResult> readHandler;

		public SyncResourceSpecification(McpSchema.Resource resource,
										 BiFunction<McpSyncServerExchange, McpSchema.ReadResourceRequest, McpSchema.ReadResourceResult> readHandler) {
			this.resource = resource;
			this.readHandler = readHandler;
		}

		public McpSchema.Resource resource() {
			return resource;
		}

		public BiFunction<McpSyncServerExchange, McpSchema.ReadResourceRequest, McpSchema.ReadResourceResult> readHandler() {
			return readHandler;
		}
	}

	/**
	 * Specification of a prompt template with its synchronous handler function.
	 */
	public static class SyncPromptSpecification {
		private final McpSchema.Prompt prompt;
		private final BiFunction<McpSyncServerExchange, McpSchema.GetPromptRequest, McpSchema.GetPromptResult> promptHandler;

		public SyncPromptSpecification(McpSchema.Prompt prompt,
									   BiFunction<McpSyncServerExchange, McpSchema.GetPromptRequest, McpSchema.GetPromptResult> promptHandler) {
			this.prompt = prompt;
			this.promptHandler = promptHandler;
		}

		public McpSchema.Prompt prompt() {
			return prompt;
		}

		public BiFunction<McpSyncServerExchange, McpSchema.GetPromptRequest, McpSchema.GetPromptResult> promptHandler() {
			return promptHandler;
		}
	}

	/**
	 * Specification of a completion handler function with synchronous execution support.
	 */
	public static class SyncCompletionSpecification {
		private final McpSchema.CompleteReference referenceKey;
		private final BiFunction<McpSyncServerExchange, McpSchema.CompleteRequest, McpSchema.CompleteResult> completionHandler;

		public SyncCompletionSpecification(McpSchema.CompleteReference referenceKey,
										   BiFunction<McpSyncServerExchange, McpSchema.CompleteRequest, McpSchema.CompleteResult> completionHandler) {
			this.referenceKey = referenceKey;
			this.completionHandler = completionHandler;
		}

		public McpSchema.CompleteReference referenceKey() {
			return referenceKey;
		}

		public BiFunction<McpSyncServerExchange, McpSchema.CompleteRequest, McpSchema.CompleteResult> completionHandler() {
			return completionHandler;
		}
	}
}